# 对数周期幂律模型的拐点预测
import math

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tushare as ts

from lppls import lppls
from datetime import datetime as dt


pro = ts.pro_api('61dc8a2e5096ba36e404d40bf123e22e756d116b0664c901f4cbe209')


def download_data():
    df = pro.index_daily(
        **{"ts_code": "000852.SH", "trade_date": "", "start_date": "", "end_date": "", "limit": "", "offset": ""},
        fields=["ts_code", "trade_date", "close", "open", "high", "low", "pre_close", "change", "pct_chg", "vol",
                "amount"]
    )
    df = df.sort_values(by='trade_date', ascending=True)

    df['Date'] = [str(s)[:10] for s in pd.to_datetime(df['trade_date'])]
    df['Open'] = df['open'].copy()
    df['High'] = df['high'].copy()
    df['Low'] = df['low'].copy()
    df['Close'] = df['close'].copy()
    df['Adj Close'] = df['close'].copy()
    df['Volume'] = df['vol'].copy()
    data = df.iloc[-2000:, :]
    return data


def train_lppl(data_kline):
    time = [pd.Timestamp.toordinal(dt.strptime(t1, '%Y-%m-%d')) for t1 in data_kline['Date']]
    price = np.log(data_kline['Adj Close'].values)
    # create observations array (expected format for LPPLS observations)
    observations = np.array([time, price])
    # set the max number for searches to perform before giving-up, the literature suggests 25
    # MAX_SEARCHES = 25
    # # instantiate a new LPPLS model with the data
    lppls_model = lppls.LPPLS(observations=observations)
    # # fit the model to the data and get back the params
    # tc, m, w, a, b, c, c1, c2, O, D = lppls_model.fit(MAX_SEARCHES)
    # lppls_model.plot_fit()
    res = lppls_model.mp_compute_nested_fits(workers=8, window_size=120, smallest_window_size=30, outer_increment=1,
                                             inner_increment=5, max_searches=25, )
    return res, lppls_model.compute_indicators(res)


def plot_bubble(res_df):
    fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(18, 10))
    ord = res_df['time'].astype('int32')
    ts = [pd.Timestamp.fromordinal(d) for d in ord]
    # plot pos bubbles
    ax1_0 = ax1.twinx()
    ax1.plot(ts, res_df['price'], color='black', linewidth=0.75)
    ax1_0.plot(ts, res_df['pos_conf'], label='bubble indicator (pos)', color='red', alpha=0.5)
    # plot neg bubbles
    ax2_0 = ax2.twinx()
    ax2.plot(ts, res_df['price'], color='black', linewidth=0.75)
    ax2_0.plot(ts, res_df['neg_conf'], label='bubble indicator (neg)', color='green', alpha=0.5)

    ax1.grid(which='major', axis='both', linestyle='--')
    ax2.grid(which='major', axis='both', linestyle='--')
    # set labels
    ax1.set_ylabel('ln(p)')
    ax2.set_ylabel('ln(p)')

    ax1_0.set_ylabel('bubble indicator (pos)')
    ax2_0.set_ylabel('bubble indicator (neg)')

    ax1_0.legend(loc=2)
    ax2_0.legend(loc=2)

    plt.xticks(rotation=45)
    return fig


if __name__ == "__main__":
    data = download_data()
    model, res_df = train_lppl(data)
    fig = plot_bubble(res_df)
    bubble_df = res_df
    bubble_df['time'] = data['trade_date'].to_list()[-len(bubble_df):]